#!/usr/bin/env python

# gsat.py
# Global Sequence Alignment Tool (V2)
# An implementation of GSAT by Vamsee Reddy, re-written
# to use the Biopython library. Much faster, more maintainable,
# more secure. Performs a Needleman-Wunsch global alignment
# of a query sequence against some number of randomly shuffled
# sequences with the same amino acid composition as the sequence
# specified as the subject; then compares the standard deviation
# of those results against the score for the query versus the subject
# for a more rigorous global alignment analysis.
#
# REQUIRES: Python 2.5 or greater, Biopython 1.47 or greater, EMBOSS 

# tl;dr:
# 1) shuffle the subject sequence X times, maintaining amino-acid comparisons
# 2) write shuffles to tempfile in fasta formats
# 3) iteratively compare query vs shuffled using Needleman-Wunsch, write score
#    to array, calculate mean and std deviation of array
# 4) compare query to subject, compare its score (x) to mean (u) over std dev (o)
#    (x-u)/o = standard score (z-score)
#
# By Erik Clarke, 2011, and Vamsee Reddy, 2010. 
# Licensed for use by academic institutions and members only.
# From the Transporter Protein Database, Saier Lab, UCSD.

import sys
import random
import re
import tempfile
import optparse
import textwrap
import numpy

from Bio.Emboss.Applications import NeedleCommandline
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_protein

class GSAT:
    '''
    Contains methods for shuffling sequences, calculating the standard
    deviations of the needle comparisons between them, and returning the
    results.
    '''
    
    def __init__(self, query, subject, shuffles, gapopen, gapextend,
    outfile, verbose=False):

        '''
        Initialize new GSAT object with passed values and some default (blank) values
        
        query, subject: strings pointing to files in FASTA format
        shuffles: int
        gapopen, gapextend: floats
        outfile: string pointing to desired outfile
        '''
        
        self.query = query
        self.subject = subject
        self.shuffles = shuffles

        self.needle = NeedleCommandline()
        self.needle.asequence = query     
        self.needle.gapopen = gapopen
        self.needle.gapextend = gapextend
        self.outfile = outfile
        self.verbose = verbose

        # initialize default values for the scores and std deviation
        self.stdev = 0
        self.zscore = 0
        self.on_score = []
        self.scores = []
        
        # initialize tempfiles
        self.shuffled_seq = tempfile.NamedTemporaryFile(mode='w+t') 
        self.shuffled_comp = tempfile.NamedTemporaryFile(mode='w+t')
        
    def __call__(self):

        '''
        Executes comparison when called; procedure spelled out in introductory
        docstring. Takes no parameters, uses existing values from __init__.
        This likely will be truncated to one method call in the future.
        '''

        Seq_fasta2 = SeqIO.read(self.subject, "fasta")
        self.shuffle(Seq_fasta2, self.shuffles)
        self.needleCompare()
        self.results()
        self.cleanup()

    def shuffle(self, subject, shuffles):

        '''
        Shuffles a given Seq object a specified number of times,
        then outputs each shuffle to a file for comparisons using
        EMBOSS needle.
        '''

        sequence = list(subject)
        for i in range(shuffles):
            random.shuffle(sequence)
            # creates a new sequence record and writes it in fasta format to /tmp/a
            shuffled_sequence = SeqRecord(Seq("".join(sequence), generic_protein), id="tmp")
            self.shuffled_seq.write(shuffled_sequence.format("fasta"))
        if self.verbose:
            print("Shuffling complete.")

    def needleCompare(self):

        '''
        Pulls the shuffled sequence tempfile and the shuffled comparison tempfile
        and does a iterative comparison of the query sequence against each of the shuffled
        sequences, adds all the scores to an array, and computes the standard deviation of
        that array. Then compares original query vs original subject and saves that comparison
        score and alignment to an outfile, if specified, or to the temp file.
        '''

        self.needle.bsequence = self.shuffled_seq.name
        self.needle.outfile = self.shuffled_comp.name
        if self.verbose:
            print("Conducting %i comparisons. This may take some time, please wait..." % self.shuffles)
        stdout, stderr = self.needle()
        self.needle.bsequence = self.subject
        self.needle.outfile = self.outfile
        stdout, stderr = self.needle()
        #rewind the files
        self.shuffled_comp.seek(0)
        

    def z_score(self):

        '''
        Computes the standard (z) score assuming on_score, scores, and
        stdev are all initiated and populated. could/should be
        rewritten to take arguments.
        '''
        
        zscore = (self.on_score[0]-self.scores.mean())/self.stdev
        zscore = round(zscore)
        self.zscore = zscore
        
        
    def find_scores(self, targetfile, scores):

        '''
        Extracts and collects the scores from a given file with one or
        more needle (or equivalent) comparisons; appends to given
        score array and returns it. This method modifies the arguments
        passed to it; specifically it populates the score[] array
        argument.

        targetfile: open handle to a given file
        scores: array to use to collect scores !This is modified
        '''

        regex = re.compile('Score: ([0-9,\.]{1,5})') # Find the number after 'Score:'
        for line in targetfile:
            tmps = regex.findall(line)
            if tmps != []: #success! add it to the scores array
                try: # but first convert it to an int or float 
                    tmps = int(tmps[0]) 
                except ValueError: #i.e., it's not an int, so try float
                    tmps = float(tmps[0])                    
                scores.append(tmps)
        return scores

    def results(self):
        '''
        Collates and (optionally) outputs the results.
        '''

        outfile = open(self.outfile)
        self.find_scores(self.shuffled_comp, self.scores)
        self.find_scores(outfile, self.on_score)
        outfile.close()
        
        # Calculate the standard deviation of the scores array
        self.scores = numpy.array(self.scores)
        self.stdev = numpy.std(self.scores)

        # Only call z_scores after stdev, scores, and on_score have
        # been populated
        self.z_score()
        
        
        if self.verbose:
            print("============ FINISHED =============")
            print("Standard score:         {0}".format(self.zscore))
            print("Query vs subject score: {0}".format(self.on_score[0]))
            print("Alignment output:       {0}".format(self.outfile))
        
        

    def cleanup(self):
        self.shuffled_seq.close()
        self.shuffled_comp.close()

####


if __name__ == '__main__':
    '''
    executes when running gsat as a standalone program to pull options
    from command line. Not really a method.
    '''

    usage = "usage: %prog query.faa subject.faa [options]"

    parser = optparse.OptionParser(
        description=textwrap.dedent('''\
        Compares two sequences to each other to determine
        a standard score (z-score) for use in homology detection.
        Scores above 10 are considered significant.'''),

        epilog=textwrap.dedent('''\
        For more info, refer to the wiki page for help at
        http://saier-144-24.ucsd.edu/groups/docs'''),
        usage = usage)

    parser.add_option('-s', action='store', default=500, type=int,
                      dest = 'shuffles',
                      help = 'number of times to shuffle subject sequence')
    parser.add_option('-o', action='store', default=8, type=int,
                      dest = 'gapopen',
                      help = 'gap-open penalty')
    parser.add_option('-e', action='store', default=2.0, type=float,
                      dest = 'gapextend',
                      help = 'gap-extend penalty')
    parser.add_option('--out', action='store', default="outfile.txt",
                      dest = 'outfile',
                      help = 'output file to write results')
    parser.add_option('-q', action='store_false', default='store_true',
                      dest = 'verbose',
                      help = 'suppress runtime messages')
    cmdline = parser.parse_args()

    # if lower than the minimum number of arguments are given, print the help and exit
    if len(cmdline[1]) < 2:
        parser.print_help()
        exit(textwrap.dedent("ERROR: First two arguments must specify query and subject fasta files."))

    args = cmdline[1]
    opts = cmdline[0]

    # instantiate our shiny new GSAT object with our hard-earned options

    newGSAT = GSAT(args[0], args[1], opts.shuffles, opts.gapopen,
                   opts.gapextend, opts.outfile, opts.verbose)
    newGSAT()
